# -*- encoding: utf-8 -*-
"""
@File    : rr.py
@Author  : lilong
@Time    : 2023/3/4 11:45 上午
"""

import numpy as np

import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Input, Dense, Layer, LSTM, Lambda, Embedding


class Attention(keras.Model):

    def __init__(self, heads, size_per_head, key_size=None, mask_right=False, **kwargs):
        """初始化模型，准备需要的类
        :param config: 参数集合
        """
        super(Attention, self).__init__(name='attention')
        self.heads = heads  # 头数量：8
        self.size_per_head = size_per_head  # 头大小：16
        self.out_dim = heads * size_per_head  # 输出维度：8x16=128
        self.key_size = key_size if key_size else size_per_head  # key大小（头大小）
        self.mask_right = mask_right

    def input_shape(self, *args, **kwargs):
        if len(args) > 0:
            inputs = args[0]
        else:
            inputs = kwargs['inputs']
        if isinstance(inputs, list):
            input_shape = [K.int_shape(x) for x in inputs]
        else:
            input_shape = K.int_shape(inputs)

        return input_shape

    def mask(self, x, mask, mode='mul'):
        """
        1-mask 存在两种结果:
        1: 0 说明该位置存在值, 相当于x对应位置的值减去0即不变
        2: 1 说明该位置对应的x处不存在值, 相当于x对应位置的值减去-1e10即负无穷
        此举的目的是计算attention score的时候对padding做mask操作
        """
        if mask is None:
            return x
        else:
            for _ in range(K.ndim(x) - K.ndim(mask)):
                mask = K.expand_dims(mask, K.ndim(mask))
            if mode == 'mul':
                return x * mask
            else:
                return x - (1 - mask) * 1e10

    def call(self, inputs, **kwargs):
        """公式是: https://pic1.zhimg.com/v2-e698e0083f4cc8d0fae45c501fb9aef8_r.jpg，
        构建模型过程,使用函数式模型，其中q指下方的q_value进行的一系列计算, k指下方的k_value的一系列计算, dk指的是self.key_size
        """

        # inputs -> [y, x, x, x_mask]
        # Attention在seq2seq中的思路是重复利用输入值与输出值作为下一个单字的输出，因此这里有两个x
        q, k, v = inputs[:3]
        v_mask, q_mask = None, None
        if len(inputs) > 3:
            v_mask = inputs[3]
            if len(inputs) > 4:
                q_mask = inputs[4]

        # 线性变换
        self.input_shape(self.key_size * self.heads, use_bias=False)
        qw = Dense(self.key_size * self.heads, input_shape=(), use_bias=False)(q)   # self.reuse(self.q_dense, q)
        kw = Dense(self.key_size * self.heads, use_bias=False)(k)   # self.reuse(self.k_dense, k)
        vw = Dense(self.out_dim, use_bias=False)(v)   # self.reuse(self.v_dense, v)

        # 形状变换
        # 确保reshape内部shape的乘积等于q_linear.shape内部的乘积即可
        qw = K.reshape(qw, (-1, K.shape(qw)[1], self.heads, self.key_size))
        # 将维度空间从(64, 128, 128)->(64, 128, 8, 16)的目的是:
        # 转化为多个低维空间最后进行拼接，形成同样维度的输出，借此丰富特性信息，降低了计算量
        # print(f"qw shape: {qw.shape}")
        kw = K.reshape(kw, (-1, K.shape(kw)[1], self.heads, self.key_size))
        vw = K.reshape(vw, (-1, K.shape(vw)[1], self.heads, self.size_per_head))

        # 维度置换
        qw = K.permute_dimensions(qw, (0, 2, 1, 3))
        # print(f"qw transpose shape: {qw.shape}")
        # 首先需要明确这里的0, 2, 1, 3并不是让qw的维度变为0, 2, 1, 3(因为也没有0这个维度)
        # 这里的0, 2, 1, 3可以认为是对维度位置对应维度的调换
        # 比如qw的维度是(64, 128, 8, 16)本身维度的位置是(0, 1, 2, 3)此时transpose希望将它的维度转变
        # 就通过它的位置信息传递, 即将其维度变为(64, 8, 128, 16)此时对应的位置就是(0, 2, 1, 3)
        kw = K.permute_dimensions(kw, (0, 2, 1, 3))
        vw = K.permute_dimensions(vw, (0, 2, 1, 3))

        # Attention
        a = tf.einsum('ijkl,ijml->ijkm', qw, kw) / self.key_size ** 0.5
        a = K.permute_dimensions(a, (0, 3, 2, 1))
        a = self.mask(a, v_mask, 'add')
        a = K.permute_dimensions(a, (0, 3, 2, 1))
        if self.mask_right:
            ones = K.ones_like(a[:1, :1])
            mask = (ones - tf.matrix_band_part(ones, -1, 0)) * 1e10
            a = a - mask
            # 以矩阵对角线(从(0, 0)处到右下角结束) num_low代表下三角, num_upper代表上三角
            # 若取负值则视作不处理, 若取正数则当做向上或者向下从起点和结束点平移对应位置画出一个三角
            # 这个三角内部的值全部填充为0
            # 此处相当于将ones的以对角线分割下半区为1上半区为0
            # 假设ones原本为:       现在则变为了:
            # [[1, 1, 1, 1]]      [[1, 0, 0, 0]]
            # [[1, 1, 1, 1]]      [[1, 1, 0, 0]]
            # [[1, 1, 1, 1]]      [[1, 1, 1, 0]]
            # [[1, 1, 1, 1]]      [[1, 1, 1, 1]]

        # 完成输出
        a = K.softmax(a)
        o = tf.einsum('ijkl,ijlm->ijkm', a, vw)
        o = K.permute_dimensions(o, (0, 2, 1, 3))
        o = K.reshape(o, (-1, K.shape(o)[1], self.out_dim))
        o = self.mask(o, q_mask, 'mul')
        return o

    def compute_output_shape(self, input_shape):
        pass


if __name__ == "__main__":
    x_sample = np.array(
        [[5, 3, 0, 0],
         [2, 2, 8, 0]])
    x_mask = Lambda(lambda x: K.cast(K.greater(K.expand_dims(x_sample, 2), 0), 'float32'))(x_sample)
    print('x_mask:', x_mask)
    # x_mask: tf.Tensor(
    # [[[1.]
    #   [1.]
    #   [0.]
    #   [0.]]
    #  [[1.]
    #   [1.]
    #   [1.]
    #   [0.]]], shape=(2, 4, 1), dtype=float32)

    # 模拟x的中间变量
    x_sample_embed = np.array(
        [
            [[1, 0, 2, 2, 4],
             [2, 2, 0, 3, 3],
             [1, 9, 2, 2, 2],
             [1, 0, 8, 2, 2]],

            [[1, 4, 2, 2, 2],
             [5, 2, 0, 3, 3],
             [1, 9, 0, 2, 2],
             [1, 2, 8, 3, 2]]
        ]
    )

    y = np.array(
        [
            [[5, 2, 0, 3, 3],
             [2, 2, 0, 3, 3],
             [1, 9, 2, 2, 2],
             [1, 2, 8, 3, 2]],

            [[1, 4, 2, 2, 2],
             [1, 0, 2, 2, 4],
             [1, 9, 0, 2, 2],
             [1, 0, 8, 2, 2]]
        ]
    )

    x = tf.cast(x_sample_embed, 'float32')
    y = tf.cast(y, 'float32')

    xy = Attention(8, 16)([y, x, x, x_mask])
    print(xy.numpy())
